
from codesecurity.feature.objects import Ast,AstNode,CommonFeatureSet
from codesecurity.tasks.code_obfuscation_detect.cod import ObfuscationDetector
import os
import hashlib

from dataclasses import dataclass,field

def calculate_file_hash(file_path, algorithm='sha256'):
    # 以二进制模式打开文件
    with open(file_path, 'rb') as f:
        # 根据传入的算法名称创建对应的哈希对象
        if algorithm == 'md5':
            hash_alg = hashlib.md5()
        elif algorithm == 'sha1':
            hash_alg = hashlib.sha1()
        elif algorithm == 'sha256':
            hash_alg = hashlib.sha256()
        elif algorithm == 'sha512':
            hash_alg = hashlib.sha512()
        else:
            raise ValueError("Unsupported algorithm")

        # 从文件中读取数据并更新哈希对象
        while True:
            data = f.read(65536)  # 每次读取 64KB 数据
            if not data:
                break
            hash_alg.update(data)

    # 返回以十六进制表示的哈希值
    return hash_alg.hexdigest()

def mini_feature(common_feature:CommonFeatureSet,name='d'):
    
    api_names=[]
    api_words=[]
    file_size=0
    file_hash=''
    entropy=0.
    
    if common_feature.property_exists(CommonFeatureSet.Addon_Api_List):
        api_names=common_feature.get_property(CommonFeatureSet.Addon_Api_List)
        
    if common_feature.property_exists(CommonFeatureSet.Addon_Common_Api_Word):
        api_words=common_feature.get_property(CommonFeatureSet.Addon_Common_Api_Word)
        
    if os.path.exists(common_feature.path):
        file_size=os.path.getsize(common_feature.path)
        entropy=ObfuscationDetector.entropy_value(common_feature.path)
        
    file_hash=calculate_file_hash(common_feature.path)
    
    return MiniCommonFeature(name,api_names,api_words,len(common_feature.ast_object.nodes),file_size,common_feature.path,file_hash,entropy)

@dataclass
class MiniCommonFeature:
    name:str
    api_names:list=field(default_factory=list)
    api_words:list=field(default_factory=list)
    node_number:int=0
    file_size:int=0
    file_path:str=""
    file_hash:str=""
    entropy:float=0.